import json
import os

json_dir = "/home/algo/geyiheng/utilgyh/allowed_seg_json"

deal_with_collect = False
deal_with_eval = False
deal_with_traverse = True

deal_with_aloha = False
deal_with_franka = True
deal_with_ur5 = True
deal_with_piper = True
set_up_piper = True

collect_name = "allowed_seg_colors-collect.json"
eval_name = "allowed_seg_colors-eval.json"
traverse_name = "allowed_seg_colors-traverse.json"
# embodiments = [
#   "aloha-agilex",
#   "franka-panda",
#   "ur5-wsg",
#   "piper"
# ]

embodiments = []


def load_embodiments():
  global embodiments, deal_with_aloha, deal_with_franka, deal_with_ur5, deal_with_piper
  if deal_with_aloha:
    embodiments.append("aloha-agilex")
  if deal_with_franka:
    embodiments.append("franka-panda")
  if deal_with_ur5:
    embodiments.append("ur5-wsg")
  if deal_with_piper:
    embodiments.append("piper")


def prepare_tasks_env(tasks):
  global deal_with_collect, deal_with_eval, deal_with_traverse, set_up_piper
  for task in tasks:
    file_path = None
    if deal_with_collect:
      file_path = os.path.join(json_dir, task, collect_name)
    if deal_with_eval:
      file_path = os.path.join(json_dir, task, eval_name)
    if deal_with_traverse:
      file_path = os.path.join(json_dir, task, traverse_name)
    if file_path is None:
      print("no file to deal with, please set deal_with_collect or deal_with_eval to True")
      return

    print(f"clearing env for {file_path}")
    clear_env(file_path)
    
    if set_up_piper:
      setup_piper(file_path)


def clear_env(file_path):  
  if not os.path.exists(file_path):
    print(f"no such file: {file_path}")
    return
  # read this json file
  # set to []: data[embodiment]["env"]["indices"] & data[embodiment]["env"]["colors"]
  with open(file_path, "r") as f:
    data = json.load(f)
  for embodiment in embodiments:
    if embodiment in data:
      if "env" in data[embodiment]:
        print(f"clearing {embodiment} env")
        data[embodiment]["env"]["indices"] = []
        data[embodiment]["env"]["colors"] = []
  with open(file_path, "w") as f:
    json.dump(data, f, indent=2)


def setup_piper(file_path):
  if not os.path.exists(file_path):
    print(f"no such file: {file_path}")
    return
  with open(file_path, "r") as f:
    data = json.load(f)
  data["piper"] = {
    "indices2obj": {
      "lhand": 99,
      "lhand_lfinger": 99,
      "lhand_rfinger": 99,
      "target": 99,
      "object": 99
    },
    "env": {
      "indices": [],
      "colors": []
    },
    "exclude": {
      "indices": [],
      "colors": []
    },
    "gripper": {
      "indices": [],
      "colors": []
    }
  }
  print(f"setting up piper in {file_path}")
  with open(file_path, "w") as f:
    json.dump(data, f, indent=2)


def main():
  load_embodiments()
  task_names = [
    "stamp_seal",
    "move_pillbottle_pad",
    "beat_block_hammer",
    "place_fan",
    "place_empty_cup",
    "rotate_qrcode",
    "stack_blocks_two",
  ]

  prepare_tasks_env(task_names)


if __name__ == "__main__":
  main()